import cv2
import numpy as np
from matplotlib import pyplot as plt

# 灰度线性变换
def hist_linear_trans(src, scale, t0):
    dst = src.astype(np.float32)
    dst = dst * scale + t0
    dst[dst > 255] = 255
    dst[dst < 0] = 0
    dst = dst.astype(np.uint8)
    return dst

# 灰度均衡化
def hist_equalize(src):
    # hist = np.zeros(256, dtype=np.float32)
    # for i in range(src.shape[0]):
    #     for j in range(src.shape[1]):
    #         hist[src[i][j]] += 1
    # hist /= src.shape[0]*src.shape[1]

    # hist = np.histogram(src, np.arange(257))[0].astype(np.float32)
    hist = cv2.calcHist([src], [0], None, [256], [0, 256]).reshape(-1)
    hist /= src.shape[0] * src.shape[1]

    lut = np.zeros(256, dtype=np.uint8)
    lut[0] = round(255 * hist[0])
    for i in range(1, 256):
        hist[i] += hist[i-1]
        h = round(255 * hist[i])
        if h > 255: h = 255
        lut[i] = h

    return cv2.LUT(src, lut)

def hist_specify(src, dst_hist):
    src_hist = cv2.calcHist([src], [0], None, [256], [0, 256]).flatten()
    src_hist /= src.shape[0] * src.shape[1]

    src_cdf = np.zeros(256, dtype=np.float32)
    dst_cdf = np.zeros(256, dtype=np.float32)
    src_cdf[0] = src_hist[0]
    dst_cdf[0] = dst_hist[0]
    for i in range(1, 256):
        src_cdf[i] = src_hist[i] + src_cdf[i-1]
        dst_cdf[i] = dst_hist[i] + dst_cdf[i-1]

    lut = np.zeros(256, dtype=np.uint8)
    for i in range(0, 256):
        # 寻求累积分布函数最接近的原图像灰度i, 和参考灰度j
        # 这等价于用目标分布函数与i所对应的累积分布函数值求差, 再寻找差的最小值做对应的j
        diff_cdf = np.abs(dst_cdf - src_cdf[i])
        j = np.argmin(diff_cdf)
        lut[i] = j

    return cv2.LUT(src, lut)

# 载入图像
filename = 'apple.bmp'
src = cv2.imread(filename, cv2.IMREAD_GRAYSCALE)
# 灰度变换
# delta = 300
# scale = (255 + 2*delta) / 255
# t0 = -delta+50
# dst = hist_linear_trans(src, scale, t0)
# dst = hist_equalize(src)
# dst = cv2.equalizeHist(src)
# 读入参考图像
# ref = cv2.imread('lena.bmp', cv2.IMREAD_GRAYSCALE)
# # 计算参考图像直方图函数
# dst_hist = cv2.calcHist([ref], [0], None, [256], [0, 255]).flatten()
# dst_hist /= ref.shape[0]*ref.shape[1]
# 产生单斜线直方图函数
dst_hist = np.zeros(256, dtype=np.float32)
k = 2.0 / (255**2)
for i in range(256):
    dst_hist[i] = k*i
# 调用直方图规定化
dst = hist_specify(src, dst_hist)

# 显示原始图像和灰度直方图
ax1 = plt.subplot(221)
ax1.set_title('Src Image')
plt.imshow(src, cmap='gray', vmin=0, vmax=255)
plt.xticks([]), plt.yticks([])
ax2 = plt.subplot(222)
ax2.set_title('Src Histogram')
plt.hist(src.ravel(), 256, [0, 256])
# 显示目标图像和灰度直方图
ax3 = plt.subplot(223)
ax3.set_title('Dst Image')
plt.imshow(dst, cmap='gray', vmin=0, vmax=255)
plt.xticks([]), plt.yticks([])
ax2 = plt.subplot(224)
ax2.set_title('Dst Histogram')
plt.hist(dst.ravel(), 256, [0, 256])

plt.show()
